﻿using System.Collections.Concurrent;
using CodeAgent.Host.Agents;
using Microsoft.AspNetCore.SignalR;
using Microsoft.SemanticKernel;
using Microsoft.SemanticKernel.ChatCompletion;

namespace CodeAgent.Host.Service;

public class CodeAgentHub(ILogger<CodeAgentHub> logger, IServiceProvider serviceProvider) : Hub
{
    private static readonly ConcurrentDictionary<string, AgentManager> Connections = new();

    public override async Task OnConnectedAsync()
    {
        var connectionId = Context.ConnectionId;

        await base.OnConnectedAsync();

        Connections.TryAdd(connectionId, new AgentManager(serviceProvider.CreateAsyncScope(), connectionId));

        logger.LogInformation("Connection {ConnectionId} connected", connectionId);
    }

    public async Task SendAgent(string message, string model)
    {
        if (Connections.TryGetValue(Context.ConnectionId, out var agent))
        {
            try
            {
                var chatHistory = new ChatMessageContent(AuthorRole.User, message);
                await agent.StartGroup(chatHistory, model);
            }
            catch (Exception e)
            {
                logger.LogError(e, "Error when sending message to agent");
            }
        }
    }

    public async Task InitManager(string token)
    {
        if (Connections.TryGetValue(Context.ConnectionId, out var agent))
        {
            await agent.Initialize(new AgentOption()
            {
                ApiKey = token,
            });

            // 通知初始成功
            await Clients.Client(Context.ConnectionId)
                .SendCoreAsync("AgentProtocol", [
                    new AgentProtocol<string>()
                    {
                        code = 200,
                        agentId = agent.AgentId,
                        message = "初始化成功",
                        agentName = "Agent",
                        type = "init",
                        result = "",
                    }
                ]);
        }
    }

    public override async Task OnDisconnectedAsync(Exception? exception)
    {
        await base.OnDisconnectedAsync(exception);

        var connectionId = Context.ConnectionId;

        if (Connections.TryRemove(connectionId, out var group))
        {
            await group.DisposeAsync();

            logger.LogInformation("Connection {ConnectionId} disposed", connectionId);
        }
    }
}